#!/usr/bin/env python
# encoding=utf-8
#codeby     道长且阻
#email      ydhcui@suliu.net/QQ664284092
#website    http://www.suliu.net

import requests
import gevent
from gevent.threadpool import ThreadPool

import socket
import importlib
from threading import Thread

try:
    from urlparse import urlsplit
except ImportError:
    from urllib.parse import urlsplit

class BaseThread(object):
    def __init__(self,queue,num_thread=10):
        self.Queue = queue
        self.num_thread = num_thread

    def todo(self,data):
        #time.sleep(0.0001)
        pass

    def run(self):
        while self.Queue.qsize() > 0:
            data = self.Queue.get()
            self.todo(data)

    def start(self,daemon=True):
        for _ in range(self.num_thread):
            t=Thread(target = self.run)
            if daemon:
                t.setDaemon(daemon)
            t.start()

class QueueThread(object):
    """队列模式的多线程，会一直等待队列中增加元素。
    任务完成后不会自动退出，直到手动调用STOP退出"""
    def __init__(self,Queue,func):
        self.Queue      = Queue
        self.func       = func
        self.threads    = []
        self.FLAG       = True
        self.STAT       = False

    def run(self):
        while self.FLAG:
            if self.STAT:
                time.sleep(0.1)
                continue
            data = self.Queue.get()
            self.func(data)

    def start(self,num_thread=10):
        for i in range(num_thread):
            self.threads.append(Thread(target = self.run))
        loops = range(len(self.threads))
        for j in loops:
            #self.threads[j].setDaemon(True)
            self.threads[j].start()
        for k in loops:
            self.threads[k].join()

    def stop(self):
        self.FLAG = False

    def pause(self):
        self.STAT = not self.STAT

    def clear(self):
        self.Queue.clear()

    @property 
    def isdone(self):
        return True and self.Queue.qsize()

class IterThread(object):
    """列表模式的多线程，列表中的任务完成后会自动退出，
    也可手动调用stop退出，这样有可能列表中的任务会没有完成"""
    def __init__(self,Queue,func):
        self.Queue      = iter(Queue)
        self.func       = func
        self.threads    = []
        self.FLAG       = True
        self.STAT       = False

    def run(self,data):
        self.func(data)

    def start(self,num_thread=10):
        while self.FLAG:
            if self.STAT:
                time.sleep(0.1)
                continue
            for i in range(num_thread):
                try:
                    data = next(self.Queue)
                except StopIteration:
                    self.FLAG = False
                    break
                self.threads.append(Thread(target = self.run, args=(data,)))
            loops = range(len(self.threads))
            for j in loops:
                #self.threads[j].setDaemon(True)
                self.threads[j].start()
            for k in loops:
                self.threads[k].join()

    def stop(self):
        self.FLAG = False

    def pause(self):
        self.STAT = not self.STAT

class GeventThread(object):
  """使用异步方式"""
  def __init__(self,threads=10):
    self.Queue   = []
    self.threads = threads
    self.timeout = 10
    self.__FLAG  = True   #stop
    self.__STAT  = False  #pause
  def recv(self,*args,**kw):
    self.handler(*args,**kw)
  def run(self):
    if self.Queue and self.handler:
      self.__FLAG  = True
      self.__STAT  = False
      pool = ThreadPool(self.threads)
      Queue = iter(self.Queue)
      while self.__FLAG:
        if self.__STAT:
          time.sleep(1)
          continue
        try:
          data = next(Queue)
          pool.spawn(self.recv, data)
        except StopIteration:
          self.__FLAG = False
          break
        gevent.wait(timeout=self.timeout)
  def stop(self):
    self.__FLAG = False
    self.__STAT = False
  def pause(self):
    self.__STAT = not self.__STAT
  def setup(self,**kwargs):
    for k,v in kwargs.items():
      setattr(self,k,v)

class Response(object):
    def __init__(self,req):
        self.req = req
    def __str__(self):
        return self.req.content 
    def __len__(self):
        return len(self.req) 

class Request(object):
    '''定义一个http请求的基类'''
    req = requests.Session()
    def __init__(self,url,data={},headers={},method='GET',version='HTTP/1.1'):
        self.url    = url
        self.data   = data
        self.headers= {
            "User-Agent":"Mozilla/5.0 (compatible; Baiduspider/2.0;)",
            "Accept-Encoding":"gzip, deflate, sdch",
            "Accept-Language":"zh-CN,zh;q=0.8",
            "Connection":"keep-alive",
            "Referer":self.url,}
        self.headers.update(headers)
        self.method = method.strip()
        self.version= version.strip()
        parser      = urlsplit(self.url)
        self.scheme = parser.scheme
        self.netloc = parser.netloc
        self.path   = parser.path
        self.query  = {}
        query       = parser.query
        if query:
           self.query = dict([q.split('=') for q in query.split('&') if '=' in q])
        if data and method=='GET':
            self.method = 'POST'

    def reload(self):
        self.url = "%s://%s%s"%(
            self.scheme,
            self.netloc,
            '%s?%s'%(self.path,
                '&'.join(['%s=%s'%(k,v) for k,v in self.query.items()]) \
                if self.query else self.path))
        return self.request()

    def request(self,timeout=None,proxy={}):
        return Response(getattr(self.req,self.method.lower())(
            self.url,data=self.data,headers=self.headers,timeout=timeout,stream=True))

    def __eq__(self,req):
        return self.path==req.path and self.query==req.query and self.data==req.data

    def __str__(self):
        s=[]
        s.append("%s %s %s"%(
            self.method.upper(),
            '%s?%s'%(self.path,
                '&'.join(['%s=%s'%(k,v) for k,v in self.query.items()])) \
                if self.query else self.path,
            self.version))
        s.append('Host: %s'%(self.netloc))
        for k,v in self.headers.items():
            s.append("%s: %s"%(k,v))
        if self.data:
            s.append('\n\n')
            s.append("&".join(["%s=%s"%(k,v) for k,v in self.data.items()]))
        s.append('\n\n')
        return '\r\n'.join(s)

class BaseDomain(object):
    '''定义一个网站的基类'''
    def __init__(self,host,port=80,domain='',ishttps=False):
        self.host         = host
        self.port         = port
        self.domain       = domain
        self.banner       = ''
        self.xpoweredby   = ''
        self.headers      = ''
        self.title        = ''
        self.status_code  = 0

        baseurl = domain if domain else '%s:%s'%(host,port) 
        self.url = 'https://%s'%baseurl if ishttps else 'http://%s'%baseurl
        self.init()

    def init(self):
        res = requests.get(self.url)
        self.banner = res.headers.get('Server')
        self.xpoweredby = res.headers.get('X-Powered-By')
        self.headers = res.headers
        self.title = ''.join(re.findall(r"<title>(.*?)</title>",res.text))

class TCPClient(object):
    def __init__(self,host,port,timeout=10):
        self.sock = socket.socket()
        self.sock.settimeout(timeout)
        self.sock.connect((host,int(port)))

    def send(self,req,chunk=1024):
        res = b''
        self.sock.sendall(req)
        while True:
            try:
                buff = self.sock.recv(chunk)
                res += buff
            except socket.error:
                break
            if not buff:
                break
        return res

    @classmethod
    def getipbyname(self,name):
        return socket.gethostbyname(name)
    @classmethod
    def getselfip(self):
        ip = '127.0.0.1'
        try:
            req = Request(url='http://1212.ip138.com/ic.asp')
            ip = ''.join(re.findall(r'\[(.*)\]',req.text))
        except:
            ip = self.getipbyname(socket.gethostname())
        finally:
            return ip

class PluginManage(object):
    plugins = {}
    def __init__(self,**kw):
        self.kwargs = kw

    def __call__(self,f):
        self.plugins[f]=self.kwargs

    @classmethod
    def test(cls,f):
        def F(self,req):
            f(self,req)
            for k,v in cls.plugins.items():
                k(req).request(**v)
        return F

    def load_plugin(self,plugin_path='./'):
        for folder in os.listdir(plugin_path):
            if '__' not in folder:
                self.plugins.append(folder)

    def add_plugin(self,plugin_name):
        if plugin_name not in self.plugins:
            if '__' not in plugin_name:
                self.plugins.append(plugin_name)

    def del_plugin(self,plugin_name):
        if plugin_name in self.plugins:
            if '__' not in plugin_name:
                self.plugins.pop(plugin_name)

#@PluginManage()
class BasePlugin(object):
    payload = ["11'\"11","<2222>"]
    def __init__(self,req):
        self.req = req
        self.settings = dict(
            name = 'BasePlugin',
            port = 80,
            thread_num = 10,
            text = '插件说明文本'
        )

    def request(self,**kwargs):
        for xss in self.payload:
            for k,v in self.req.query.items():
                self.req.query[k]=xss
                if xss in self.req.reload().text:
                   self.handler(self.req,xss)
                   return

    def handler(self,req,xss):
        log.error("XSS %s[payload]%s"%(req.url,xss))

class ConnectionError(requests.ConnectionError):pass
class BaseScan(object):pass

